from langchain.memory import ConversationBufferMemory
from langchain.chains import ConversationChain

class ChatMemorySystem:
    def __init__(self, llm):
        self.memory = ConversationBufferMemory(
            memory_key="chat_history",  # 存储对话历史的字段名
            return_messages=True  # 返回对话历史    
        )
        
        # 创建对话链
        self.chain = ConversationChain(
            llm=llm,
            memory=self.memory,
            verbose=True
        )

    def get_memory(self):
        return self.memory

    def clear_memory(self):
        self.memory = []
